import time
import math
import pygame
import numpy as np
import argparse

from board import board
from model import loadAgent
from pygame_2048.utils import set_window, make_move, ARROW_KEYS, draw_button
from pygame_2048.board import Board


def main():
    # AI agent plays 2048
    parser = argparse.ArgumentParser()
    parser.add_argument('--load_ckpt', type=str, required=True)
    parser.add_argument('--speed', type=float, default=0.01, help="secs to sleep after each step. smaller for faster playing")
    args = parser.parse_args()

    # init pygame
    pygame.init()
    pygame.font.init()
    clock = pygame.time.Clock()

    window = set_window()

    pygame_board = Board()
    pygame_board.draw(window)
    pygame.display.update()

    # init agent and game
    agent = loadAgent(args.load_ckpt)

    input('Press Enter to Start')     # wait for start
    while True:
        for _ in pygame.event.get():
            pass        # 不加这个for循环pygame窗口会卡死

        tiles = [int(math.log(i, 2)) if i else 0 for line in pygame_board.board for i in line]
        game = board(tiles)
        next_games = [game.left(), game.right(), game.up(), game.down()]
        # print(game)
        # for i, (g, score) in enumerate(next_games):
        #     print(i)
        #     print(g, score)
        state_values = agent.evaulate(next_games)
        state_values = [state_values[i] if (np.array(next_game.tile) - np.array(game.tile)).any()
                        else min(state_values)-100000 for i, (next_game, _) in enumerate(next_games)]
        # print(state_values)
        action = np.argmax(state_values)
        # pygame_action = (action + 2) % 4
        time.sleep(args.speed)
        make_move(pygame_board, ARROW_KEYS[action], window)

        if pygame_board.no_moves():
            print('game over, closing in 10 secs')
            time.sleep(10)
            break

if __name__ == '__main__':
    main()